package com.ruoyi.ratelimit;

import com.google.common.util.concurrent.RateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Aspect // 声明这是一个切面类
@Component
public class RateLimitAspect {

    /**
     * 使用ConcurrentHashMap缓存每个方法的限流器。
     * Key: 方法名，Value: RateLimiter实例
     */
    private final ConcurrentHashMap<String, RateLimiter> limiterMap = new ConcurrentHashMap<>();

    /**
     * 环绕通知，拦截所有被 @RateLimit 注解的方法
     */
    @Around("@annotation(rateLimitAnnotation)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimitAnnotation) throws Throwable {
        // 1. 获取注解参数
        double  permitsPerSecond = rateLimitAnnotation.value();
        long timeout = rateLimitAnnotation.timeout();
        TimeUnit timeUnit = rateLimitAnnotation.timeUnit();

        // 2. 获取方法唯一标识（例如：类名+方法名）
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        String methodKey = signature.getDeclaringTypeName() + "." + signature.getName();

        // 3. 获取或创建该方法的限流器
        RateLimiter rateLimiter = limiterMap.computeIfAbsent(methodKey,
                key -> RateLimiter.create(permitsPerSecond)
        );

        // 4. 尝试获取令牌
        boolean acquired;
        if (timeout > 0) {
            // 在指定的超时时间内尝试获取令牌
            acquired = rateLimiter.tryAcquire(timeout, timeUnit);
        } else {
            // 非阻塞方式，获取不到立即返回false
            acquired = rateLimiter.tryAcquire();
        }

        // 5. 根据获取结果决定是否放行
        if (acquired) {
            // 获取成功，执行原方法
            return joinPoint.proceed();
        } else {
            // 获取失败，触发限流处理（例如：返回错误信息）
            // 这里可以根据需要返回JSON、抛出异常等
            throw new RuntimeException("请求过于频繁，请稍后再试");
            // 或者返回一个友好的Result对象：
            // return Result.error(429, "请求过于频繁，请稍后再试");
        }
    }
}
